/**
 * Copyright 2020 Huawei Technologies Co., Ltd
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#include "tf_util.h"

namespace aicpu {
static uint64_t g_kernel_id = 0;

uint64_t GenerateUniqueId() {
  if (g_kernel_id == ULLONG_MAX) {
    g_kernel_id = 0;
  }
  return g_kernel_id++;
}

TFDataType ConvertGeDataType2TfDataType(ge::DataType data_type, bool is_ref) {
  static const std::map<ge::DataType, TFDataType> data_type_map = {
      {ge::DataType::DT_FLOAT16, TFDataType::DT_HALF},
      {ge::DataType::DT_FLOAT, TFDataType::DT_FLOAT},
      {ge::DataType::DT_DOUBLE, TFDataType::DT_DOUBLE},
      {ge::DataType::DT_INT8, TFDataType::DT_INT8},
      {ge::DataType::DT_UINT8, TFDataType::DT_UINT8},
      {ge::DataType::DT_INT16, TFDataType::DT_INT16},
      {ge::DataType::DT_UINT16, TFDataType::DT_UINT16},
      {ge::DataType::DT_INT32, TFDataType::DT_INT32},
      {ge::DataType::DT_UINT32, TFDataType::DT_UINT32},
      {ge::DataType::DT_INT64, TFDataType::DT_INT64},
      {ge::DataType::DT_UINT64, TFDataType::DT_UINT64},
      {ge::DataType::DT_BOOL, TFDataType::DT_BOOL},
      {ge::DataType::DT_RESOURCE, TFDataType::DT_RESOURCE},
      {ge::DataType::DT_STRING, TFDataType::DT_STRING},
      {ge::DataType::DT_STRING_REF, TFDataType::DT_STRING_REF},
      {ge::DataType::DT_COMPLEX64, TFDataType::DT_COMPLEX64},
      {ge::DataType::DT_COMPLEX128, TFDataType::DT_COMPLEX128},
      {ge::DataType::DT_QINT8, TFDataType::DT_QINT8},
      {ge::DataType::DT_QUINT8, TFDataType::DT_QUINT8},
      {ge::DataType::DT_QINT16, TFDataType::DT_QINT16},
      {ge::DataType::DT_QUINT16, TFDataType::DT_QUINT16},
      {ge::DataType::DT_QINT32, TFDataType::DT_QINT32},
      {ge::DataType::DT_VARIANT, TFDataType::DT_VARIANT}};

  const auto diff = TFDataType::DT_FLOAT_REF - TFDataType::DT_FLOAT;
  std::map<ge::DataType, TFDataType>::const_iterator iter = data_type_map.find(data_type);
  if (iter != data_type_map.end()) {
    if (is_ref) {
      TFDataType tf_data_type = iter->second;
      // Avoid mishandling DT_STRING_REF
      if (tf_data_type < diff) {
        tf_data_type = static_cast<TFDataType>(tf_data_type + diff);
      }
      return tf_data_type;
    } else {
      return iter->second;
    }
  } else {
    return TFDataType::DT_INVALID;
  }
}

bool ConvertString2TfDataType(const string &elem, TFDataType &data_type) {
  static const map<string, TFDataType> data_type_map = {
      {"DT_INT8", TFDataType::DT_INT8},
      {"DT_UINT8", TFDataType::DT_UINT8},
      {"DT_INT16", TFDataType::DT_INT16},
      {"DT_UINT16", TFDataType::DT_UINT16},
      {"DT_INT32", TFDataType::DT_INT32},
      {"DT_INT64", TFDataType::DT_INT64},
      {"DT_UINT32", TFDataType::DT_UINT32},
      {"DT_UINT64", TFDataType::DT_UINT64},
      {"DT_FLOAT16", TFDataType::DT_HALF},
      {"DT_FLOAT", TFDataType::DT_FLOAT},
      {"DT_DOUBLE", TFDataType::DT_DOUBLE},
      {"DT_BOOL", TFDataType::DT_BOOL},
      {"DT_COMPLEX64", TFDataType::DT_COMPLEX64},
      {"DT_COMPLEX128", TFDataType::DT_COMPLEX128},
      {"DT_STRING", TFDataType::DT_STRING},
      {"DT_STRING_REF", TFDataType::DT_STRING_REF},
      {"DT_RESOURCE", TFDataType::DT_RESOURCE},
      {"DT_QINT8", TFDataType::DT_QINT8},
      {"DT_QUINT8", TFDataType::DT_QUINT8},
      {"DT_QINT16", TFDataType::DT_QINT16},
      {"DT_QUINT16", TFDataType::DT_QUINT16},
      {"DT_QINT32", TFDataType::DT_QINT32},
      {"DT_VARIANT", TFDataType::DT_VARIANT}};

  auto it = data_type_map.find(elem);
  // only data_type in configure file exists in data_type_map, save it and return
  if (it != data_type_map.end()) {
      data_type = it->second;
      return true;
  }
  return false;
}
} // namespace aicpu
